package org.cupk2.service.controller;

import com.cupk2.lucene.Cupk2LuceneFace;
import com.cupk2.lucene.document.CupkDocument;
import com.cupk2.lucene.dto.IndexInfoDto;
import com.cupk2.lucene.inter.Cupk2LuceneIndexInter;
import com.cupk2.lucene.utils.CupkLuceneUtils;
import org.cupk2.auth.dto.KeyValueDto;
import org.cupk2.service.event.utils.KnowIndexDocumentBuilder;
import org.cupk2.service.lucene.doc.KnowIndexDocument;
import org.cupk2.zhzm.enums.KnowFileTypeEnum;
import org.cupk2.zhzm.service.ZhzmKnowServiceInter;
import org.cupk2.tools.db.DataQuery;
import org.cupk2.tools.db.DataResult;
import org.cupk2.tools.db.ResultDataHandle;
import org.cupk2.tools.files.Cupk2FileUtils;
import org.cupk2.tools.web.CupkResponseCode;
import org.cupk2.tools.web.CupkResponseResult;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 索引管理
 *
 * @author cbtg自动生成  2025-1-2 22:01:57
 */
@RestController
@RequestMapping("/api/lucene")
public class WebLuceneController {
    @Autowired
    private ZhzmKnowServiceInter zhzmKnowServiceImpl;
    @Autowired
    private KnowIndexDocumentBuilder snowIndexDocumentBuilder;
    /**
     *
     * @return
     */
    @PreAuthorize("@cupkAction.has('lucene.info')")
    @PostMapping("/info")
    public CupkResponseResult info() {
        Map<String, IndexInfoDto> infos = new HashMap<>();
        IndexInfoDto textIndex = new IndexInfoDto();
        textIndex.setTitle("文本索引");
        textIndex.setSize(CupkLuceneUtils.getIndexSize(CupkLuceneUtils.getLuceneIndexFilePath()));
        textIndex.setPath(CupkLuceneUtils.getLuceneIndexFilePath());
        infos.put("text", textIndex);
        IndexInfoDto vectorIndex = new IndexInfoDto();
        vectorIndex.setTitle("向量索引");
        vectorIndex.setSize(CupkLuceneUtils.getIndexSize(CupkLuceneUtils.getLuceneEmbeddingFilePath()));
        vectorIndex.setPath(CupkLuceneUtils.getLuceneEmbeddingFilePath());
        infos.put("vector", vectorIndex);
        return CupkResponseResult.getInstance(CupkResponseCode.SUCESS, infos);
    }


    @PreAuthorize("@cupkAction.has('lucene.index')")
    @PostMapping("/index")
    public CupkResponseResult index(@RequestBody KeyValueDto key) throws IOException {
        if (key.getKey().equals("text")) {
            new Thread(() -> {
                reIndexAllKnow();
            }).start();
        }
        if (key.getKey().equals("vector")) {
            new Thread(() -> {
                reVectorAllKnow();
            }).start();
        }
        return CupkResponseResult.getInstance(CupkResponseCode.SUCESS);
    }


    /**
     * 重新创建所有知识的向量索引
     */
    private void reVectorAllKnow() {
        {//创建知识索引
            DataQuery query = new DataQuery();
            query.setPageSize(100);
            int page = 1;
            Cupk2LuceneIndexInter indexs = Cupk2LuceneFace.getIndexService(CupkLuceneUtils.getLuceneEmbeddingFilePath());
            while (true) {
                List<CupkDocument> docs = new ArrayList<>();
                query.setPage(page);
                page++;
                DataResult result = zhzmKnowServiceImpl.searchZhzmKnow(query);
                if (result.getData().size() <= 0) {
                    break;
                } else {
                    result.runDataHandle(new ResultDataHandle() {
                        @Override
                        public void handle(Map<String, Object> row) {
                            String knowid = (String) row.get("ID");
                            List<String> files = zhzmKnowServiceImpl.getKnowFiles(knowid, null);
                            for (String fileid : files) {
                                docs.addAll(snowIndexDocumentBuilder.getVectorDoc(knowid, fileid));
                            }
                        }
                    });
                }
                indexs.createIndex(docs);
            }
        }
    }


    /**
     * 重新创建所有知识的索引
     */
    private void reIndexAllKnow() {
        {//创建知识索引
            DataQuery query = new DataQuery();
            query.setPageSize(100);
            int page = 1;
            Cupk2LuceneIndexInter indexs = Cupk2LuceneFace.getIndexService();
            while (true) {
                List<CupkDocument> docs = new ArrayList<>();
                query.setPage(page);
                page++;
                DataResult result = zhzmKnowServiceImpl.searchZhzmKnow(query);
                if (result.getData().size() <= 0) {
                    break;
                } else {
                    result.runDataHandle(new ResultDataHandle() {
                        @Override
                        public void handle(Map<String, Object> row) {
                            String knowid = (String) row.get("ID");
                            KnowIndexDocument doc = snowIndexDocumentBuilder.getKnowDoc(knowid);
                            if (doc != null) {
                                docs.add(doc.getDoc());
                                List<String> files = zhzmKnowServiceImpl.getKnowFiles(knowid, KnowFileTypeEnum.KNOW_FILE.getKey());
                                for (String fileid : files) {
                                    try {
                                        KnowIndexDocument filedoc = snowIndexDocumentBuilder.getFileDoc(knowid, fileid);
                                        if (filedoc != null) {
                                            docs.add(filedoc.getDoc());
                                        }
                                    } catch (IOException e) {
                                        e.printStackTrace();
                                    }
                                }
                            }
                        }
                    });
                }
                indexs.createIndex(docs);
            }
        }
    }


    @PreAuthorize("@cupkAction.has('lucene.del')")
    @PostMapping("/del")
    public CupkResponseResult del(@RequestBody KeyValueDto key) throws IOException {
        if (key.getKey().equals("text")) {

            if (Path.of(CupkLuceneUtils.getLuceneIndexFilePath()).toFile().exists()) {
                Cupk2FileUtils.deleteFolderRecursively(Path.of(CupkLuceneUtils.getLuceneIndexFilePath()));
            }
        }

        if (key.getKey().equals("vector")) {

            if (Path.of(CupkLuceneUtils.getLuceneIndexFilePath()).toFile().exists()) {
                Cupk2FileUtils.deleteFolderRecursively(Path.of(CupkLuceneUtils.getLuceneEmbeddingFilePath()));
            }
        }
        return CupkResponseResult.getInstance(CupkResponseCode.SUCESS);
    }
}
